import streamlit as st
from typing import Dict, Any, List

def init_session_state():
    """初始化Streamlit会话状态"""
    
    # 文件上传相关
    if 'uploaded_files' not in st.session_state:
        st.session_state.uploaded_files = []
    
    if 'file_info' not in st.session_state:
        st.session_state.file_info = {}
    
    # 确保widget相关的键不被预设，让Streamlit自动管理
    # 这些键包括: page_selector, chunking_strategy 等
    
    # 处理配置相关
    if 'cleaning_config' not in st.session_state:
        st.session_state.cleaning_config = {
            'remove_empty_lines': True,
            'remove_duplicates': True,
            'remove_headers_footers': False,
            'remove_special_chars': True,
            'normalize_text': True,
            'fix_line_breaks': True
        }
    
    if 'chunking_config' not in st.session_state:
        st.session_state.chunking_config = {
            'strategy': '固定长度分段',
            'chunk_size': 500,
            'overlap_size': 50,
            'custom_separator': '\n\n',
            'max_depth': 3,
            'preserve_structure': True
        }
    
    # 处理状态相关
    if 'processing_status' not in st.session_state:
        st.session_state.processing_status = 'idle'
    
    # 处理结果相关
    if 'processed_documents' not in st.session_state:
        st.session_state.processed_documents = []
    
    if 'processing_history' not in st.session_state:
        st.session_state.processing_history = []
    
    # 预览相关
    if 'preview_data' not in st.session_state:
        st.session_state.preview_data = None
    
    if 'current_preview_page' not in st.session_state:
        st.session_state.current_preview_page = 0

def update_processing_status(status: str):
    """更新处理状态"""
    st.session_state.processing_status = status

def add_processed_document(doc_data: Dict[str, Any]):
    """添加处理完成的文档"""
    st.session_state.processed_documents.append(doc_data)
    
    # 添加到历史记录
    history_item = {
        'timestamp': doc_data.get('processing_time'),
        'filename': doc_data.get('metadata', {}).get('title', 'Unknown'),
        'chunks_count': doc_data.get('metadata', {}).get('total_chunks', 0),
        'strategy': st.session_state.chunking_config['strategy'],
        'status': 'completed'
    }
    st.session_state.processing_history.append(history_item)

def clear_session_data():
    """清除会话数据"""
    st.session_state.uploaded_files = []
    st.session_state.file_info = {}
    st.session_state.processed_documents = []
    st.session_state.preview_data = None
    st.session_state.processing_status = 'idle'

def get_file_info(filename: str) -> Dict[str, Any]:
    """获取文件信息"""
    return st.session_state.file_info.get(filename, {})

def set_file_info(filename: str, info: Dict[str, Any]):
    """设置文件信息"""
    st.session_state.file_info[filename] = info